# ui/heatmap_viewer.py

from PyQt5.QtWidgets import QWidget, QVBoxLayout, QLabel, QPushButton, QFileDialog
from PyQt5.QtGui import QPixmap
import os


class HeatmapViewer(QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.init_ui()

    def init_ui(self, Qt=None):
        layout = QVBoxLayout(self)

        self.heatmap_label = QLabel("请选择热图文件", self)
        self.heatmap_label.setAlignment(Qt.AlignCenter)
        layout.addWidget(self.heatmap_label)

        select_btn = QPushButton("选择热图", self)
        select_btn.clicked.connect(self.select_heatmap)
        layout.addWidget(select_btn)

        back_btn = QPushButton("返回主菜单", self)
        back_btn.clicked.connect(lambda: self.parent().switch_page(0))
        layout.addWidget(back_btn)

        self.setLayout(layout)

    def select_heatmap(self):
        path, _ = QFileDialog.getOpenFileName(self, "选择热图文件", "results/", "Image Files (*.png *.jpg *.svg)")
        if path:
            pixmap = QPixmap(path)
            self.heatmap_label.setPixmap(pixmap.scaledToWidth(600))

    def show_heatmap(self, heatmap_file="results/figures/community_similarity.png"):
        if os.path.exists(heatmap_file):
            pixmap = QPixmap(heatmap_file)
            self.heatmap_label.setPixmap(pixmap.scaledToWidth(600))
        else:
            self.heatmap_label.setText("热图不存在，请先生成")
